package com.tool.grab.vpn.task;

import android.net.VpnService;

import com.tool.grab.vpn.CacheMap;
import com.tool.grab.vpn.header.UDPPacket;
import com.tool.grab.vpn.tunnel.TunnelInterface;
import com.tool.grab.vpn.tunnel.UDPTunnel;

import java.io.IOException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.util.Iterator;
import java.util.Set;

public class UDPService implements Runnable {
    private VpnService service;
    private Selector selector;
    private boolean isRun;

    private final CacheMap<Integer,UDPTunnel> udpConnection = new CacheMap<>(20, new CacheMap.DeleteCallback<UDPTunnel>() {
        @Override
        public void delete(UDPTunnel udpTunnel) {
            udpTunnel.onClose();
        }
    });

    public UDPService(VpnService service){
        this.service = service;
        try {
            selector = Selector.open();
        }catch (IOException e){
            stop();
        }
    }

    public void start(){
        Thread thread = new Thread(this,"UDPService");
        thread.start();
    }

    public void processPacket(int port,UDPPacket packet){
        UDPTunnel tunnel = udpConnection.get(port);

        if(tunnel == null){
            tunnel = new UDPTunnel(selector);
            try {
                tunnel.connect(service,packet);
                udpConnection.put(port,tunnel);
            } catch (IOException e) {}

        }else{
            tunnel.processPacket(packet);
        }
    }

    @Override
    public void run() {
        isRun = true;
        try {
            while (isRun){
                int select = selector.select();

                if(select == 0){
                    Thread.sleep(10);
                    continue;
                }

                Set<SelectionKey> keySet = selector.selectedKeys();

                if(keySet == null) continue;

                Iterator<SelectionKey> keys = keySet.iterator();

                while (keys.hasNext()){
                    SelectionKey key = keys.next();
                    if(key.isValid()){
                        Object object = key.attachment();
                        if(object instanceof TunnelInterface){
                            ((TunnelInterface) object).keyReady(key);
                        }
                    }
                }
                keys.remove();
            }
        } catch (Exception e) {
            this.stop();
        }

    }

    private void stop(){
        isRun = false;
        if(selector != null){
            try {
                selector.close();
            } catch (IOException e) {}finally {
                selector = null;
            }
        }

        for (Integer i : udpConnection.keySet()){
            udpConnection.get(i).onClose();
        }

        udpConnection.clear();
    }
}
